#include "com_seetaface2_SeetaFace2JNI.h"
#include "JniUtils.h"
#include "CJFieldStruct.h"

#include <string>

#include <seeta/FaceDetector2.h>
#include <seeta/PointDetector2.h>
#include <seeta/FaceRecognizer.h>
#include <seeta/FaceCropper2.h>

using namespace std;

#define LOGD printf //debug日志
#define LOGE printf //error日志

//是否初始化成功
static bool is_initialized = false;

seeta::FaceDetector2* FD = NULL;
seeta::PointDetector2* PD = NULL;
seeta::FaceRecognizer2* FR = NULL;

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    initModel
* Signature: (Ljava/lang/String;)Z
*/
JNIEXPORT jboolean JNICALL Java_com_seetaface2_SeetaFace2JNI_initModel
(JNIEnv *env, jobject obj, jstring model_dir) {
	//如果已初始化则直接返回
	if (is_initialized) {
		return true;
	}

	if (NULL == model_dir) {
		return false;
	}
	string dir = env->GetStringUTFChars(model_dir, 0);
	if (dir.empty()) {
		return false;
	}
	string last_char = dir.substr(dir.length() - 1, 1);
	if ("\\" == last_char) {
		dir = dir.substr(0, dir.length() - 1) + "/";
	}
	else if (last_char != "/") {
		//目录补齐/
		dir += "/";
	}

	//检测模型路径
	string detect_model_path = dir + "SeetaFaceDetector2.0.ats";
	//对齐模型路径
	string align_model_path = dir + "SeetaPointDetector2.0.pts5.ats";
	//识别模型路径
	string recognizer_model_path = dir + "SeetaFaceRecognizer2.0.ats";

	//3个模型文件都要存在，否则初始化失败
	FILE* fp = fopen(detect_model_path.c_str(), "r");
	if (fp == NULL) {
		LOGE("detect model file[%s] not exist!\n", detect_model_path.c_str());
		return false;
	}
	fclose(fp);

	fp = fopen(align_model_path.c_str(), "r");
	if (fp == NULL) {
		LOGE("align model file[%s] not exist!\n", align_model_path.c_str());
		return false;
	}
	fclose(fp);

	fp = fopen(recognizer_model_path.c_str(), "r");
	if (fp == NULL) {
		LOGE("recognizer model file[%s] not exist!\n", recognizer_model_path.c_str());
		return false;
	}
	fclose(fp);

	// Initialize face detection model
	LOGD("Initialize face detection model: %s\n", detect_model_path.c_str());
	FD = new seeta::FaceDetector2(detect_model_path.c_str());

	// Initialize face alignment model 
	LOGD("Initialize face alignment model: %s\n", align_model_path.c_str());
	PD = new seeta::PointDetector2(align_model_path.c_str());

	// Initialize face Identification model 
	LOGD("Initialize face recognizer model: %s\n", recognizer_model_path.c_str());
	FR = new seeta::FaceRecognizer2(recognizer_model_path.c_str());

	LOGD("Initialized successfully!\n");
	is_initialized = true;
	return true;
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    detect
* Signature: (Lcom/seetaface2/model/SeetaImageData;)[Lcom/seetaface2/model/SeetaRect;
*/
JNIEXPORT jobjectArray JNICALL Java_com_seetaface2_SeetaFace2JNI_detect__Lcom_seetaface2_model_SeetaImageData_2
(JNIEnv *env, jobject obj, jobject jimg) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return NULL;
	}
	//将java的SeetaImageData对象转成C++的SeetaImageData
	SeetaImageData image = JniUtils::toSeetaImageData(env, jimg);

	int face_num;    //保存人脸数量
	SeetaRect *face = FD->Detect(image, &face_num);//检测人脸

	//LOGD("detect face num: %i", face_num);
	if (face_num <= 0) {//没有检测到人脸
		env->DeleteLocalRef(jimg);
		return NULL;
	}

	//获取SeetaRect java对象
	RectFieldID rect_field = JniUtils::getRectFieldID(env);
	
	//构造一个指向jrect_class类的一维数组对象，该对象数组初始大小为face_num
	jclass clazz = env->FindClass(rect_field.className);
	jobjectArray jrect_array = env->NewObjectArray(face_num, clazz, NULL);
	for (int i = 0; i < face_num; ++i, ++face) {
		//创建java对象: new SeetaRect()
		jobject jrect = env->NewObject(clazz, rect_field.constructor);
		//设置人脸范围数据
		env->SetIntField(jrect, rect_field.x, face->x);//x坐标
		env->SetIntField(jrect, rect_field.y, face->y);//y坐标
		env->SetIntField(jrect, rect_field.width, face->width);//宽度
		env->SetIntField(jrect, rect_field.height, face->height);//高度
		//添加到对象数组
		env->SetObjectArrayElement(jrect_array, i, jrect);
		//删除局部引用
		env->DeleteLocalRef(jrect);
	}
	env->DeleteLocalRef(jimg);

	return jrect_array;
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    detect
* Signature: (Lcom/seetaface2/model/SeetaImageData;[Lcom/seetaface2/model/SeetaRect;)[Lcom/seetaface2/model/SeetaPointF;
*/
JNIEXPORT jobjectArray JNICALL Java_com_seetaface2_SeetaFace2JNI_detect__Lcom_seetaface2_model_SeetaImageData_2_3Lcom_seetaface2_model_SeetaRect_2
(JNIEnv *env, jobject obj, jobject jimg, jobjectArray jfaces) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return NULL;
	}
	//将java的SeetaImageData对象转成C++的SeetaImageData
	SeetaImageData image = JniUtils::toSeetaImageData(env, jimg);
	//将java的SeetaRect[]对象转成C++的SeetaRect[]
	SeetaRect *face = JniUtils::toSeetaRectArray(env, jfaces);
	SeetaPointF *points = PD->Detect(image, *face);
	if (!points) {
		env->DeleteLocalRef(jimg);
		env->DeleteLocalRef(jfaces);
		return NULL;
	}
	//获取SeetaPointF java对象
	PointFieldID point_field = JniUtils::getPointFieldID(env);
	
	int landmark_num = PD->LandmarkNum();
	//构造一个指向jpoint_class类的一维数组对象，该对象数组初始大小为landmark_num
	jclass clazz = env->FindClass(point_field.className);
	jobjectArray jpoint_array = env->NewObjectArray(landmark_num, clazz, NULL);
	for (int i = 0; i < landmark_num; ++i, ++points) {
		//创建java对象: new SeetaPointF()
		jobject jpoint = env->NewObject(clazz, point_field.constructor); 
		//设置特征坐标数据
		env->SetDoubleField(jpoint, point_field.x, points->x);//x坐标
		env->SetDoubleField(jpoint, point_field.y, points->y);//y坐标
		//添加到对象数组
		env->SetObjectArrayElement(jpoint_array, i, jpoint);
		//删除局部引用
		env->DeleteLocalRef(jpoint);
	}

	//delete[] image.data;
	//delete[] face;
	//delete[] points;
	env->DeleteLocalRef(jimg);
	env->DeleteLocalRef(jfaces);

	return jpoint_array;
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    compare
* Signature: (Lcom/seetaface2/model/SeetaImageData;Lcom/seetaface2/model/SeetaImageData;)F
*/
JNIEXPORT jfloat JNICALL Java_com_seetaface2_SeetaFace2JNI_compare
(JNIEnv *env, jobject obj, jobject jimg1, jobject jimg2) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return -1;
	}
	const int landmark_num = 5;

	//将java的SeetaImageData对象转成C++的SeetaImageData
	SeetaImageData image1 = JniUtils::toSeetaImageData(env, jimg1);
	//LOGD("width=%i,height=%i,channel=%i,data=%x\n", image1.width, image1.height, image1.channels, image1.data);
	SeetaRect *face1 = FD->Detect(image1);
	if (!face1) {
		//LOGD("No face detected in image1.\n");
		return -1.1f;
	}
	SeetaPointF points1[landmark_num];
	PD->DetectEx(image1, *face1, points1);

	SeetaImageData image2 = JniUtils::toSeetaImageData(env, jimg2);
	SeetaRect *face2 = FD->Detect(image2);
	if (!face2) {
		//LOGD("No face detected in image2.\n");
		return -1.2f;
	}
	SeetaPointF points2[landmark_num];
	PD->DetectEx(image2, *face2, points2);

	float similar = FR->Compare(image1, points1, image2, points2);
	
	//delete[] image1.data;
	//delete[] image2.data;
	//delete face1;
	//delete face2;
	//delete[] points1;
	//delete[] points2;
	env->DeleteLocalRef(jimg1);
	env->DeleteLocalRef(jimg2);

	return similar;
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    register
* Signature: (Lcom/seetaface2/model/SeetaImageData;)I
*/
JNIEXPORT jint JNICALL Java_com_seetaface2_SeetaFace2JNI_register
(JNIEnv *env, jobject obj, jobject jimg) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return -1;
	}
	SeetaImageData image = JniUtils::toSeetaImageData(env, jimg);
	SeetaRect *face = FD->Detect(image);
	if (!face) {
		//LOGD("No face detected.\n");
		return -1;
	}
	SeetaPointF *points = PD->Detect(image, *face);
	env->DeleteLocalRef(jimg);

	if (!points) {
		return -1;
	}
	return FR->Register(image, points);  // Reture -1 if failed.
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    recognize
* Signature: (Lcom/seetaface2/model/SeetaImageData;)Lcom/seetaface2/model/RecognizeResult;
*/
JNIEXPORT jobject JNICALL Java_com_seetaface2_SeetaFace2JNI_recognize
(JNIEnv *env, jobject obj, jobject jimg) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return NULL;
	}
	SeetaImageData image = JniUtils::toSeetaImageData(env, jimg);

	SeetaRect *face = FD->Detect(image);
	if (!face) {
		//LOGD("No face detected.\n");
		return NULL;
	}
	SeetaPointF *points = PD->Detect(image, *face);

	float similar = 0;
	int index = FR->Recognize(image, points, &similar);
	//printf("recognize: index=%i, similar=%f", index, similar);

	//获取SeetaRect java对象
	RecognizeFieldID result_field = JniUtils::getRecognizeFieldID(env);
	//创建java对象: new RecognizeResult()
	jobject jrecognize = JniUtils::newObject(env, result_field);
	//设置特征坐标数据
	env->SetIntField(jrecognize, result_field.index, index);//数据库索引
	env->SetFloatField(jrecognize, result_field.similar, similar);//相似度

	env->DeleteLocalRef(jimg);
	return jrecognize;
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    clear
* Signature: ()V
*/
JNIEXPORT void JNICALL Java_com_seetaface2_SeetaFace2JNI_clear
(JNIEnv *env, jobject obj) {
	FR->Clear();
}

/*
* Class:     com_seetaface2_SeetaFace2JNI
* Method:    crop
* Signature: (Lcom/seetaface2/model/SeetaImageData;)[B
*/
JNIEXPORT jbyteArray JNICALL Java_com_seetaface2_SeetaFace2JNI_crop
(JNIEnv *env, jobject obj, jobject jimg) {
	if (!is_initialized) {
		LOGE("model not initialized! Please use initModel() to init.\n");
		return NULL;
	}
	SeetaImageData image = JniUtils::toSeetaImageData(env, jimg);
	//LOGD("---width=%i, height=%i, channel=%i, size=%i\n", image.width, image.height, image.channels, strlen((char *)image.data));
	SeetaRect *face = FD->Detect(image);
	if (!face) {
		//LOGD("No face detected.\n");
		return NULL;
	}
	seeta::FaceCropper2* FC = new seeta::FaceCropper2();
	SeetaPointF *points = PD->Detect(image, *face);
	SeetaImageData cropped_face = FC->Crop(image, points);

	//LOGD("width=%i, height=%i, channel=%i\n", cropped_face.width, cropped_face.height, cropped_face.channels);
	
	char *buf = (char*)cropped_face.data;
	int size = 256 * 256 * 3;
	jbyteArray data = env->NewByteArray(size);
	env->SetByteArrayRegion(data, 0, size, (jbyte *)buf);
	//LOGD("data=%s, data.size=%i\n", buf, size);
	return data;
}